package cn.amebius.simpledocker.websocket;

import cn.amebius.simpledocker.config.DockerFactory;
import cn.amebius.simpledocker.entity.model.DockerClientModel;
import cn.amebius.simpledocker.thread.AttachConsoleThread;
import com.github.dockerjava.api.DockerClient;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import javax.annotation.PostConstruct;
import java.io.IOException;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author 阿伟
 * @date 2023/9/6 19:00
 */
@Slf4j
@Component
public class ContainerAttachWebSocketHandler extends TextWebSocketHandler {
    private static DockerFactory dockerFactory;

    private static ConcurrentHashMap<WebSocketSession,PipedOutputStream > userMap;

    private static ConcurrentHashMap<PipedOutputStream,AttachConsoleThread> threadMap;

    @Autowired
    public void setService(DockerFactory dockerFactory) {
        ContainerAttachWebSocketHandler.dockerFactory = dockerFactory;
        ContainerAttachWebSocketHandler.userMap = new ConcurrentHashMap<>();
        ContainerAttachWebSocketHandler.threadMap = new ConcurrentHashMap<>();
    }
    @PostConstruct
    public void printWsStartInfo() {
        log.info("容器操作WebSocket初始化完成！");
    }

    // 连接建立时触发
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws IOException {
        PipedInputStream pipedInputStream = new PipedInputStream();
        PipedOutputStream pipedOutputStream = new PipedOutputStream();
        pipedInputStream.connect(pipedOutputStream);
        userMap.put(session,pipedOutputStream);
        log.info("【ContainerAttachWs】连接:{}",session.getRemoteAddress());
        log.info("【ContainerAttachWs】连接数：{}", ContainerAttachWebSocketHandler.userMap.size());
        Integer key = (Integer) session.getAttributes().get("key");
        String containerId = (String) session.getAttributes().get("containerId");
        DockerClientModel dcm = dockerFactory.getByKey(key);
        DockerClient dc = dcm.getDockerClient();
        AttachConsoleThread act = new AttachConsoleThread(session,pipedInputStream,dc,containerId);
        threadMap.put(pipedOutputStream,act);
        act.start();

    }
    // 处理收到的文本消息
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        PipedOutputStream inputStream = userMap.get(session);
        String msg = message.getPayload() + "\r";
        inputStream.write(msg.getBytes());
        inputStream.flush();
    }

    // 连接关闭时触发
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws IOException {
        PipedOutputStream pipedOutputStream = userMap.get(session);
        pipedOutputStream.close();
        userMap.remove(session);
        AttachConsoleThread act = threadMap.get(pipedOutputStream);
        act.interrupt();
        threadMap.remove(pipedOutputStream);
        log.info("【ContainerAttachWs】退出：{}",session.getRemoteAddress());
        log.info("【ContainerAttachWs】连接数：{}", userMap.size());
    }
}
